import math
import numpy as np
from matplotlib import pyplot

def anint(x, lx):
    res=x
    if x>0.5*lx:
        res-=lx
    elif x<-0.5*lx:
        res+=lx
    return res

def distancia(lbox,datos,i,j):
    dis2=0
    drx=datos[i][0]-datos[j][0]
    drx=anint(drx, lbox)
    dry=datos[i][1]-datos[j][1]
    dry=anint(dry, lbox)
    drz=datos[i][2]-datos[j][2]
    drz=anint(drz, lbox)
    dis2=drx*drx+dry*dry+drz*drz
    return  [math.sqrt(dis2),drx,dry,drz]

def load(name):
    datos=[]
    for line in file(name):
        line=line.replace('\n','')
        val=line.split(' ')
        res=[]
        for v in val:
            if v!='':
                res.append(float(v))
        datos.append(res)
    return datos

def eKin(datos):
    ek=0.0;
    np=0
    for dt in datos:
        np+=1
        vx=dt[3]
        vy=dt[4]
        vz=dt[5]
        v2=vx*vx+vy*vy+vz*vz
        ek+=v2
    return 0.5*ek

def force(rcut, lbox, datos):
    forc=np.zeros((len(datos),3))  
    for i in range(len(datos)-1):
        for j in range(i+1,len(datos)):
            [dis,dx,dy,dz] = distancia(lbox,datos,i,j)
            if dis<rcut:
                f=24.*(2.*math.pow(dis,-13.0)-math.pow(dis,-7.0))
                fx=f*dx/dis
                fy=f*dy/dis
                fz=f*dz/dis           
            
                forc[i][0]+=fx
                forc[i][1]+=fy
                forc[i][2]+=fz
            
                forc[j][0]-=fx
                forc[j][1]-=fy
                forc[j][2]-=fz
    return forc

def force2(rcut, lbox, datos):
    forc=np.zeros((len(datos),3))  
    for i in range(len(datos)):
        for j in range(len(datos)):
            if i!=j:
                [dis,dx,dy,dz] = distancia(lbox,datos,i,j)
                if dis<rcut:
                    f=24.*(2.*math.pow(dis,-13.0)-math.pow(dis,-7.0))
                    fx=f*dx/dis
                    fy=f*dy/dis
                    fz=f*dz/dis        
                
                    forc[i][0]+=fx
                    forc[i][1]+=fy
                    forc[i][2]+=fz
    return forc

def Epoten(rcut,lbox,datos):
    ePot=0.
    for i in range(len(datos)):
        for j in range(len(datos)):
            if i!=j:
                res= distancia(lbox,datos,i,j)
                dis=res[0]     
                potdm=0.           
                if dis<rcut:
                    dis3=dis*dis*dis
                    dis6=dis3*dis3
                    dis12=dis6*dis6
                    potdm = (1./dis12-1./dis6)
                    #print i,j,dis,potdm
                ePot+=potdm
    return 2.*ePot


def testEnergy():
    lbox=5.0    
    xnew=[]
    ynew=[]
    for i in range(10):
        name='/home/jonk/codeEspresso/test/posvelforce'+str(i)+'.dat'
        datos = load(name)
        Ek=eKin(datos)
        Temp=Ek*2./(3.*300)
        rcut=2.0
        potLJcut = 300.*4.*(math.pow(rcut,-12.0)-math.pow(rcut,-6.0))
        ePot=Epoten(rcut,lbox,datos)
        xnew.append(Ek)
        ynew.append(ePot)
        #print i,Ek, ePot
    
    name='/home/jonk/codeEspresso/test/energia.dat'
    datos = load(name)
    x=[]
    y=[]
    for dt in datos:
        x.append(dt[0])
        y.append(dt[1])
    
    pyplot.plot(x,y,'ro')
    pyplot.plot(xnew,ynew,'bx')
    pyplot.show()


#Para comparar con las fuerzas de "Espresso" es necesario poner el dumping del termostato a 0
#para obtener unicamente las fuerzas debidas al potencial

#testEnergy()
name='/home/jonk/codeEspresso/test/posvelforce1.dat'
rcut=math.pow(2.,1./6.)
lbox=13.572088083
datos = load(name)
forcesEspresso=[]
for dt in datos:
    forcesEspresso.append([dt[6],dt[7],dt[8]])
#print forcesEspresso[0]
forceMio = force2(rcut,lbox,datos)
#print forceMio[0]
x=[]
y1=[]
y2=[]
for i in range(len(forceMio)):
    x.append(i)
    y1.append(forcesEspresso[i][0])
    y2.append(forceMio[i][0])
print y1
print y2
pyplot.plot(x,y1,'ro')
pyplot.plot(x,y2,'bx')
pyplot.show()

